package io.wen.bd.s8m1;

import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.AllWindowedStream;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingProcessingTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class Task2_FlinkWindow {
    public static void main(String[] args) throws Exception {
        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        final DataStream<String> logStream = env.addSource(new LogGenerateSourceFunction());

        final KeyedStream<String, Integer> keyedLogStream = logStream.keyBy(new KeySelector<String, Integer>() {
            @Override
            public Integer getKey(String value) throws Exception {
                return (int) (value.charAt(value.length() - 1));
            }
        });

        final SingleOutputStreamOperator<Integer> keyedStatStream = keyedLogStream
                .timeWindow(Time.seconds(5))
//                .timeWindow(Time.seconds(5), Time.seconds(1))
                .aggregate(new StatWarnAggregationFunction())
                .flatMap(new FlatMapFunction<Map<Integer, Integer>, Integer>() {
                    @Override
                    public void flatMap(Map<Integer, Integer> value, Collector<Integer> out) throws Exception {
                        for (Integer key : value.keySet()) {
                            // 统计窗口内次数大于3
                            if (value.get(key) > 3) {
                                out.collect(key);
                            }
                        }
                    }
                });

        keyedStatStream.print();

        env.execute();
    }


    /**
     *  日志生成 SourceFunction
     *
     */
    public static class LogGenerateSourceFunction implements SourceFunction<String> {

        private boolean running = true;

        @Override
        public void run(SourceContext<String> sourceContext) throws Exception {
            Random random = new Random();

            while (running) {
                SimpleDateFormat df = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS");
                // 20%概率生成WARN日志, body值为0到9，用作分区key
                int levelRand = random.nextInt(5);
                String level = levelRand == 0 ? "WARN" : "INFO";
                int value = random.nextInt(10);
                String msg = String.format("%s [%s] %s", df.format(new Date()), level, value);
                sourceContext.collect(msg);
                Thread.sleep(20 + random.nextInt(10));
            }
        }

        @Override
        public void cancel() {
            running = false;
        }
    }

    /**
     * 日志统计聚合函数
     */
    public static class StatWarnAggregationFunction implements AggregateFunction<String, Map<Integer, Integer>, Map<Integer, Integer>> {

        private Pattern p = Pattern.compile("\\[WARN\\] (\\d+)");
        @Override
        public Map<Integer, Integer> createAccumulator() {
            return new HashMap<>();
        }

        @Override
        public Map<Integer, Integer> add(String value, Map<Integer, Integer> acc) {
            Matcher matcher = p.matcher(value);
            if (matcher.find()) {
                int v = Integer.parseInt(matcher.group(1));
                if (!acc.containsKey(v)) {
                    acc.put(v, 1);
                } else {
                    acc.put(v, acc.get(v) + 1);
                }
            }
            return acc;
        }

        @Override
        public Map<Integer, Integer> getResult(Map<Integer, Integer> acc) {
            return acc;
        }

        @Override
        public Map<Integer, Integer> merge(Map<Integer, Integer> a, Map<Integer, Integer> b) {
            for (int key: b.keySet()) {
                if (!a.containsKey(key)) {
                    a.put(key, b.get(key));
                } else {
                    a.put(key, a.get(key) + b.get(key));
                }
            }
            return a;
        }
    }
}


